
package com.cu.machinelearning.neural;

import java.util.*;

/**
 * ResNetDemo 类实现了简化版的 ResNet 神经网络结构。
 * 包含多个残差块、全连接输出层，并支持前向传播、反向传播、损失计算与准确率评估功能。
 */
public class ResNetDemo {
    // 网络配置参数
    private final int inputSize;        // 输入特征维度
    private final int numClasses;       // 输出类别数
    private final int batchSize;        // 批处理大小
    private final double learningRate;  // 学习率
    private final List<ResidualBlock> blocks;  // 残差块列表
    private final DenseLayer finalLayer;      // 最终输出层

    // 随机数生成器
    private final Random random;

    /**
     * 构造函数，初始化 ResNet 网络结构。
     *
     * @param inputSize     输入特征维度
     * @param numClasses    分类数量（输出维度）
     * @param batchSize     批次大小
     * @param learningRate  学习率
     * @param numFilters    每个残差块使用的滤波器数量序列
     */
    public ResNetDemo(int inputSize, int numClasses, int batchSize, double learningRate,
                      int... numFilters) {
        this.inputSize = inputSize;
        this.numClasses = numClasses;
        this.batchSize = batchSize;
        this.learningRate = learningRate;
        this.random = new Random(42);  // 固定种子，保证可复现性
        this.blocks = new ArrayList<>();

        // 创建残差块
        int currentSize = inputSize;
        for (int filters : numFilters) {
            blocks.add(new ResidualBlock(currentSize, filters, batchSize, random));
            currentSize = filters;
        }

        // 创建最终输出层
        this.finalLayer = new DenseLayer(currentSize, numClasses, batchSize, random);
    }

    /**
     * 前向传播过程：依次通过所有残差块和最终输出层。
     *
     * @param input 输入数据 (batchSize x inputSize)
     * @return 网络输出 (batchSize x numClasses)，已应用 softmax
     */
    public double[][] forward(double[][] input) {
        double[][] current = input;

        // 经过所有残差块
        for (ResidualBlock block : blocks) {
            current = block.forward(current);
        }

        // 全局平均池化 (简化处理，实际实现可更复杂)
        double[][] pooled = new double[batchSize][current[0].length];
        for (int i = 0; i < batchSize; i++) {
            System.arraycopy(current[i], 0, pooled[i], 0, current[i].length);
        }

        // 最终输出层
        return finalLayer.forward(pooled);
    }

    /**
     * 反向传播过程：从输出层开始逐层回传梯度。
     *
     * @param outputGradient 输出梯度 (batchSize x numClasses)
     */
    public void backward(double[][] outputGradient) {
        // 从输出层开始反向传播
        double[][] gradient = finalLayer.backward(outputGradient, learningRate);

        // 反向传播经过所有残差块
        for (int i = blocks.size() - 1; i >= 0; i--) {
            gradient = blocks.get(i).backward(gradient, learningRate);
        }
    }

    /**
     * 使用交叉熵损失函数计算预测结果与真实标签之间的损失。
     *
     * @param predictions 预测概率分布 (batchSize x numClasses)
     * @param labels      真实标签（one-hot 编码）(batchSize x numClasses)
     * @return 平均损失值
     */
    public double computeLoss(double[][] predictions, double[][] labels) {
        double totalLoss = 0.0;
        for (int i = 0; i < batchSize; i++) {
            for (int j = 0; j < numClasses; j++) {
                totalLoss -= labels[i][j] * Math.log(Math.max(predictions[i][j], 1e-10));
            }
        }
        return totalLoss / batchSize;
    }

    /**
     * 计算模型在当前批次上的分类准确率。
     *
     * @param predictions 预测概率分布 (batchSize x numClasses)
     * @param labels      真实标签（one-hot 编码）(batchSize x numClasses)
     * @return 准确率（介于 0 到 1 之间）
     */
    public double computeAccuracy(double[][] predictions, double[][] labels) {
        int correct = 0;
        for (int i = 0; i < batchSize; i++) {
            int predClass = argMax(predictions[i]);
            int trueClass = argMax(labels[i]);
            if (predClass == trueClass) {
                correct++;
            }
        }
        return (double) correct / batchSize;
    }

    /**
     * 辅助函数：找到数组中最大值的索引位置。
     *
     * @param array 待查找的一维数组
     * @return 最大元素的下标索引
     */
    private int argMax(double[] array) {
        int maxIndex = 0;
        for (int i = 1; i < array.length; i++) {
            if (array[i] > array[maxIndex]) {
                maxIndex = i;
            }
        }
        return maxIndex;
    }

    /**
     * 残差块类，包含两个卷积层、批量归一化以及可能存在的捷径连接。
     */
    private static class ResidualBlock {
        private final ConvLayer conv1;
        private final BatchNorm bn1;
        private final ConvLayer conv2;
        private final BatchNorm bn2;
        private final ConvLayer shortcut;  // 用于维度不匹配时的捷径连接
        private double[][] inputCache;     // 缓存输入用于反向传播

        /**
         * 构造函数，构建一个残差块结构。
         *
         * @param inputChannels   输入通道数
         * @param outputChannels  输出通道数
         * @param batchSize       批次大小
         * @param random          随机数生成器
         */
        public ResidualBlock(int inputChannels, int outputChannels, int batchSize, Random random) {
            // 第一个卷积层
            this.conv1 = new ConvLayer(inputChannels, outputChannels, 3, 1, 1, batchSize, random);
            this.bn1 = new BatchNorm(outputChannels, batchSize);

            // 第二个卷积层
            this.conv2 = new ConvLayer(outputChannels, outputChannels, 3, 1, 1, batchSize, random);
            this.bn2 = new BatchNorm(outputChannels, batchSize);

            // 如果输入输出维度不同，创建1x1卷积用于捷径连接
            if (inputChannels != outputChannels) {
                this.shortcut = new ConvLayer(inputChannels, outputChannels, 1, 1, 0, batchSize, random);
            } else {
                this.shortcut = null;
            }
        }

        /**
         * 前向传播操作：执行主路径和捷径连接的组合。
         *
         * @param input 当前层输入数据
         * @return 处理后的输出数据
         */
        public double[][] forward(double[][] input) {
            this.inputCache = input;  // 缓存输入

            // 主路径
            double[][] x = conv1.forward(input);
            x = bn1.forward(x, true);  // 训练模式
            x = relu(x);
            x = conv2.forward(x);
            x = bn2.forward(x, true);  // 训练模式

            // 捷径连接
            double[][] shortcut = (this.shortcut != null) ? this.shortcut.forward(input) : input;

            // 残差连接：主路径 + 捷径
            return add(x, shortcut);
        }

        /**
         * 反向传播操作：将梯度分别传递给主路径和捷径分支。
         *
         * @param gradient       上一层传入的梯度
         * @param learningRate   当前学习率
         * @return 返回到上一层的梯度
         */
        public double[][] backward(double[][] gradient, double learningRate) {
            // 残差梯度分配
            double[][] gradMain = gradient.clone();
            double[][] gradShortcut = gradient.clone();

            // 主路径反向传播
            gradMain = bn2.backward(gradMain, learningRate);
            gradMain = conv2.backward(gradMain, learningRate);
            gradMain = reluDerivative(gradMain, conv2.getOutputCache());
            gradMain = bn1.backward(gradMain, learningRate);
            gradMain = conv1.backward(gradMain, learningRate);

            // 捷径路径反向传播
            if (shortcut != null) {
                gradShortcut = shortcut.backward(gradShortcut, learningRate);
            }

            // 合并梯度
            return add(gradMain, gradShortcut);
        }

        /**
         * ReLU 激活函数：对输入矩阵进行逐元素非线性变换。
         *
         * @param x 输入矩阵
         * @return 应用了 ReLU 的输出矩阵
         */
        private double[][] relu(double[][] x) {
            double[][] result = new double[x.length][x[0].length];
            for (int i = 0; i < x.length; i++) {
                for (int j = 0; j < x[i].length; j++) {
                    result[i][j] = Math.max(0, x[i][j]);
                }
            }
            return result;
        }

        /**
         * ReLU 导数函数：根据原始输入计算其对应的梯度掩码。
         *
         * @param gradient 梯度矩阵
         * @param x        对应的原始输入缓存
         * @return 应用了 ReLU 导数的梯度矩阵
         */
        private double[][] reluDerivative(double[][] gradient, double[][] x) {
            double[][] result = new double[gradient.length][gradient[0].length];
            for (int i = 0; i < gradient.length; i++) {
                for (int j = 0; j < gradient[i].length; j++) {
                    result[i][j] = gradient[i][j] * (x[i][j] > 0 ? 1 : 0);
                }
            }
            return result;
        }

        /**
         * 矩阵加法运算：将两个相同形状的二维数组按位相加。
         *
         * @param a 第一个矩阵
         * @param b 第二个矩阵
         * @return 相加后的新矩阵
         */
        private double[][] add(double[][] a, double[][] b) {
            double[][] result = new double[a.length][a[0].length];
            for (int i = 0; i < a.length; i++) {
                for (int j = 0; j < a[i].length; j++) {
                    result[i][j] = a[i][j] + b[i][j];
                }
            }
            return result;
        }
    }

    /**
     * 卷积层类，模拟标准卷积操作。
     */
    private static class ConvLayer {
        /**
         * 卷积层类，用于执行卷积神经网络中的卷积操作
         *
         * @param inputChannels 输入通道数
         * @param outputChannels 输出通道数
         * @param kernelSize 卷积核大小
         * @param stride 步长
         * @param padding 填充大小
         * @param batchSize 批处理大小
         * @param kernels 卷积核权重数组，维度为[输出通道数][输入通道数][卷积核大小×卷积核大小]
         * @param biases 偏置数组，长度为输出通道数
         * @param inputCache 输入数据缓存
         * @param outputCache 输出数据缓存
         */
        private final int inputChannels;
        private final int outputChannels;
        private final int kernelSize;
        private final int stride;
        private final int padding;
        private final int batchSize;
        private final double[][][] kernels;  // [output][input][kernelSize*kernelSize]
        private final double[] biases;
        private double[][] inputCache;
        private double[][] outputCache;

        /**
         * 构造函数，初始化卷积核和偏置项。
         *
         * @param inputChannels   输入通道数
         * @param outputChannels  输出通道数
         * @param kernelSize      卷积核尺寸
         * @param stride          步长
         * @param padding         填充大小
         * @param batchSize       批次大小
         * @param random          随机数生成器
         */
        public ConvLayer(int inputChannels, int outputChannels, int kernelSize,
                         int stride, int padding, int batchSize, Random random) {
            this.inputChannels = inputChannels;
            this.outputChannels = outputChannels;
            this.kernelSize = kernelSize;
            this.stride = stride;
            this.padding = padding;
            this.batchSize = batchSize;

            // 初始化权重 (Xavier初始化)
            double std = Math.sqrt(2.0 / (inputChannels * kernelSize * kernelSize));
            this.kernels = new double[outputChannels][inputChannels][kernelSize * kernelSize];
            for (int o = 0; o < outputChannels; o++) {
                for (int i = 0; i < inputChannels; i++) {
                    for (int k = 0; k < kernelSize * kernelSize; k++) {
                        kernels[o][i][k] = random.nextGaussian() * std;
                    }
                }
            }

            // 初始化偏置
            this.biases = new double[outputChannels];
            Arrays.fill(biases, 0.01);
        }

        /**
         * 前向传播操作：执行卷积运算。
         *
         * @param input 输入数据
         * @return 卷积输出
         */
        public double[][] forward(double[][] input) {
            this.inputCache = input;

            // 简化实现：假设输入为展平的特征图
            int outputSize = outputChannels;
            double[][] output = new double[batchSize][outputSize];

            // 卷积计算
            for (int b = 0; b < batchSize; b++) {
                for (int o = 0; o < outputChannels; o++) {
                    double sum = biases[o];
                    for (int i = 0; i < inputChannels; i++) {
                        for (int k = 0; k < kernelSize * kernelSize; k++) {
                            sum += input[b][i * kernelSize * kernelSize + k] * kernels[o][i][k];
                        }
                    }
                    output[b][o] = sum;
                }
            }

            this.outputCache = output;
            return output;
        }

        /**
         * 反向传播操作：更新权重并计算输入梯度。
         *
         * @param gradient       来自上一层的梯度
         * @param learningRate   当前学习率
         * @return 输入梯度
         */
        public double[][] backward(double[][] gradient, double learningRate) {
            int inputSize = inputChannels * kernelSize * kernelSize;
            double[][] inputGradient = new double[batchSize][inputSize];
            double[][][] kernelGradient = new double[outputChannels][inputChannels][kernelSize * kernelSize];
            double[] biasGradient = new double[outputChannels];

            // 计算梯度
            for (int b = 0; b < batchSize; b++) {
                for (int o = 0; o < outputChannels; o++) {
                    biasGradient[o] += gradient[b][o] / batchSize;

                    for (int i = 0; i < inputChannels; i++) {
                        for (int k = 0; k < kernelSize * kernelSize; k++) {
                            int inputIndex = i * kernelSize * kernelSize + k;
                            kernelGradient[o][i][k] += inputCache[b][inputIndex] * gradient[b][o] / batchSize;
                            inputGradient[b][inputIndex] += kernels[o][i][k] * gradient[b][o];
                        }
                    }
                }
            }

            // 更新权重和偏置
            for (int o = 0; o < outputChannels; o++) {
                biases[o] -= learningRate * biasGradient[o];
                for (int i = 0; i < inputChannels; i++) {
                    for (int k = 0; k < kernelSize * kernelSize; k++) {
                        kernels[o][i][k] -= learningRate * kernelGradient[o][i][k];
                    }
                }
            }

            return inputGradient;
        }

        /**
         * 获取输出缓存，供后续层使用。
         *
         * @return 输出缓存
         */
        public double[][] getOutputCache() {
            return outputCache;
        }
    }

    /**
     * 批量归一化层类，用于加速训练并提高稳定性。
     */
    private static class BatchNorm {
        private final int channels;
        private final int batchSize;
        private final double[] gamma;  // 缩放参数
        private final double[] beta;   // 偏移参数
        private final double eps = 1e-5;
        private final double momentum = 0.9;

        // 运行时统计量（用于推理）
        private double[] runningMean;
        private double[] runningVar;

        // 缓存用于反向传播
        private double[][] xNormCache;
        private double[] meanCache;
        private double[] varCache;

        /**
         * 构造函数，初始化缩放和平移参数及运行时统计量。
         *
         * @param channels   特征通道数
         * @param batchSize  批次大小
         */
        public BatchNorm(int channels, int batchSize) {
            this.channels = channels;
            this.batchSize = batchSize;

            // 初始化参数
            this.gamma = new double[channels];
            Arrays.fill(gamma, 1.0);
            this.beta = new double[channels];
            Arrays.fill(beta, 0.0);

            // 初始化运行时统计量
            this.runningMean = new double[channels];
            this.runningVar = new double[channels];
            Arrays.fill(runningVar, 1.0);
        }

        /**
         * 前向传播操作：对输入进行标准化处理。
         *
         * @param input       输入数据
         * @param isTraining  是否处于训练状态
         * @return 标准化后的输出
         */
        public double[][] forward(double[][] input, boolean isTraining) {
            double[][] output = new double[batchSize][channels];

            if (isTraining) {
                // 计算批次均值和方差
                meanCache = new double[channels];
                varCache = new double[channels];

                // 计算均值
                for (int c = 0; c < channels; c++) {
                    double sum = 0.0;
                    for (int b = 0; b < batchSize; b++) {
                        sum += input[b][c];
                    }
                    meanCache[c] = sum / batchSize;
                }

                // 计算方差
                for (int c = 0; c < channels; c++) {
                    double sum = 0.0;
                    for (int b = 0; b < batchSize; b++) {
                        sum += Math.pow(input[b][c] - meanCache[c], 2);
                    }
                    varCache[c] = sum / batchSize + eps;
                }

                // 归一化
                xNormCache = new double[batchSize][channels];
                for (int b = 0; b < batchSize; b++) {
                    for (int c = 0; c < channels; c++) {
                        xNormCache[b][c] = (input[b][c] - meanCache[c]) / Math.sqrt(varCache[c]);
                        output[b][c] = gamma[c] * xNormCache[b][c] + beta[c];
                    }
                }

                // 更新运行时统计量
                for (int c = 0; c < channels; c++) {
                    runningMean[c] = momentum * runningMean[c] + (1 - momentum) * meanCache[c];
                    runningVar[c] = momentum * runningVar[c] + (1 - momentum) * varCache[c];
                }
            } else {
                // 推理模式，使用运行时统计量
                for (int b = 0; b < batchSize; b++) {
                    for (int c = 0; c < channels; c++) {
                        double xNorm = (input[b][c] - runningMean[c]) / Math.sqrt(runningVar[c] + eps);
                        output[b][c] = gamma[c] * xNorm + beta[c];
                    }
                }
            }

            return output;
        }

        /**
         * 反向传播操作：计算梯度并更新参数。
         *
         * @param gradient       输入梯度
         * @param learningRate   当前学习率
         * @return 返回到上一层的梯度
         */
        public double[][] backward(double[][] gradient, double learningRate) {
            double[][] inputGradient = new double[batchSize][channels];
            double[] gammaGradient = new double[channels];
            double[] betaGradient = new double[channels];

            // 计算beta和gamma的梯度
            for (int c = 0; c < channels; c++) {
                for (int b = 0; b < batchSize; b++) {
                    betaGradient[c] += gradient[b][c] / batchSize;
                    gammaGradient[c] += gradient[b][c] * xNormCache[b][c] / batchSize;
                }
            }

            // 计算输入的梯度
            for (int b = 0; b < batchSize; b++) {
                for (int c = 0; c < channels; c++) {
                    double term1 = gradient[b][c] * gamma[c];
                    double term2 = 1.0 / Math.sqrt(varCache[c]);
                    double term3 = (1.0 / batchSize) * (1 - 1.0 / Math.sqrt(varCache[c]) *
                            Math.pow(inputGradient[b][c] - meanCache[c], 2) / varCache[c]);

                    inputGradient[b][c] = term1 * (term2 - term3);
                }
            }

            // 更新参数
            for (int c = 0; c < channels; c++) {
                gamma[c] -= learningRate * gammaGradient[c];
                beta[c] -= learningRate * betaGradient[c];
            }

            return inputGradient;
        }
    }

    /**
     * 全连接层类，用于最后的分类任务。
     */
    private static class DenseLayer {
        private final int inputSize;
        private final int outputSize;
        private final int batchSize;
        private final double[][] weights;  // [output][input]
        private final double[] biases;
        private double[][] inputCache;

        /**
         * 构造函数，初始化权重和偏置。
         *
         * @param inputSize    输入维度
         * @param outputSize   输出维度
         * @param batchSize    批次大小
         * @param random       随机数生成器
         */
        public DenseLayer(int inputSize, int outputSize, int batchSize, Random random) {
            this.inputSize = inputSize;
            this.outputSize = outputSize;
            this.batchSize = batchSize;

            // 初始化权重 (Xavier初始化)
            double std = Math.sqrt(2.0 / (inputSize + outputSize));
            this.weights = new double[outputSize][inputSize];
            for (int o = 0; o < outputSize; o++) {
                for (int i = 0; i < inputSize; i++) {
                    weights[o][i] = random.nextGaussian() * std;
                }
            }

            // 初始化偏置
            this.biases = new double[outputSize];
            Arrays.fill(biases, 0.01);
        }

        /**
         * 前向传播操作：执行线性变换并应用 softmax。
         *
         * @param input 输入数据
         * @return softmax 输出
         */
        public double[][] forward(double[][] input) {
            this.inputCache = input;
            double[][] output = new double[batchSize][outputSize];

            // 矩阵乘法 + 偏置
            for (int b = 0; b < batchSize; b++) {
                for (int o = 0; o < outputSize; o++) {
                    double sum = biases[o];
                    for (int i = 0; i < inputSize; i++) {
                        sum += input[b][i] * weights[o][i];
                    }
                    output[b][o] = softmax(sum, output, b, o);  // 应用softmax
                }
            }

            return output;
        }

        /**
         * Softmax 函数：确保输出是一个合法的概率分布。
         *
         * @param x       当前神经元的加权和
         * @param output  整体输出矩阵
         * @param batch   当前样本编号
         * @param index   当前神经元编号
         * @return softmax 归一化的值
         */
        private double softmax(double x, double[][] output, int batch, int index) {
            // 先找到最大值，防止指数溢出
            double maxVal = Double.NEGATIVE_INFINITY;
            for (int o = 0; o < outputSize; o++) {
                double val = biases[o];
                for (int i = 0; i < inputSize; i++) {
                    val += inputCache[batch][i] * weights[o][i];
                }
                if (val > maxVal) {
                    maxVal = val;
                }
            }

            // 计算指数和
            double sumExp = 0.0;
            for (int o = 0; o < outputSize; o++) {
                double val = biases[o];
                for (int i = 0; i < inputSize; i++) {
                    val += inputCache[batch][i] * weights[o][i];
                }
                sumExp += Math.exp(val - maxVal);
            }

            return Math.exp(x - maxVal) / sumExp;
        }

        /**
         * 反向传播操作：计算梯度并更新权重和偏置。
         *
         * @param gradient       输入梯度
         * @param learningRate   当前学习率
         * @return 返回到上一层的梯度
         */
        public double[][] backward(double[][] gradient, double learningRate) {
            double[][] inputGradient = new double[batchSize][inputSize];
            double[][] weightGradient = new double[outputSize][inputSize];
            double[] biasGradient = new double[outputSize];

            // 计算梯度
            for (int b = 0; b < batchSize; b++) {
                for (int o = 0; o < outputSize; o++) {
                    biasGradient[o] += gradient[b][o] / batchSize;

                    for (int i = 0; i < inputSize; i++) {
                        weightGradient[o][i] += inputCache[b][i] * gradient[b][o] / batchSize;
                        inputGradient[b][i] += weights[o][i] * gradient[b][o];
                    }
                }
            }

            // 更新权重和偏置
            for (int o = 0; o < outputSize; o++) {
                biases[o] -= learningRate * biasGradient[o];
                for (int i = 0; i < inputSize; i++) {
                    weights[o][i] -= learningRate * weightGradient[o][i];
                }
            }

            return inputGradient;
        }
    }

    // 示例用法
    public static void main(String[] args) {
        // 配置网络参数
        int inputSize = 32 * 32 * 3;  // 例如CIFAR-10图像大小
        int numClasses = 10;
        int batchSize = 32;
        double learningRate = 0.001;

        // 创建ResNet网络 (3个残差块，分别有64, 128, 256个滤波器)
        ResNetDemo resnet = new ResNetDemo(inputSize, numClasses, batchSize, learningRate, 64, 128, 256);

        // 生成随机测试数据
        Random random = new Random();
        double[][] input = new double[batchSize][inputSize];
        double[][] labels = new double[batchSize][numClasses];

        for (int i = 0; i < batchSize; i++) {
            // 随机输入
            for (int j = 0; j < inputSize; j++) {
                input[i][j] = random.nextGaussian() * 0.1;
            }

            // 随机标签 (one-hot)
            int label = random.nextInt(numClasses);
            Arrays.fill(labels[i], 0.0);
            labels[i][label] = 1.0;
        }

        // 训练一个批次
        for (int epoch = 0; epoch < 10; epoch++) {
            // 前向传播
            double[][] predictions = resnet.forward(input);

            // 计算损失和准确率
            double loss = resnet.computeLoss(predictions, labels);
            double accuracy = resnet.computeAccuracy(predictions, labels);

            System.out.printf("Epoch %d: Loss = %.4f, Accuracy = %.4f%n", epoch, loss, accuracy);

            // 计算输出梯度 (交叉熵损失的导数)
            double[][] outputGradient = new double[batchSize][numClasses];
            for (int i = 0; i < batchSize; i++) {
                for (int j = 0; j < numClasses; j++) {
                    outputGradient[i][j] = (predictions[i][j] - labels[i][j]) / batchSize;
                }
            }

            // 反向传播
            resnet.backward(outputGradient);
        }
    }
}
